import os
import torch
from torchvision.datasets import ImageFolder
from torchvision import transforms
from torch.utils.data import DataLoader, Dataset, random_split
import matplotlib.pyplot as plt

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

device = "cuda" if torch.cuda.is_available() else "cpu"

data_path = "./dataset"

# Better data augmentations for training set
train_transform = transforms.Compose([
    transforms.RandomResizedCrop(224, scale=(0.8, 1.0)),
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(15),
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
    transforms.TrivialAugmentWide(num_magnitude_bins=6),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    transforms.RandomErasing(p=0.2, scale=(0.02, 0.1), ratio=(0.3, 3.3)),
])

# Transform for validation and test sets
val_test_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# Load the full dataset without transforms to get the image paths and labels
full_dataset = ImageFolder(root=data_path)

# Split the dataset into training and validation sets
train_size = int(0.8 * len(full_dataset))
val_size = len(full_dataset) - train_size
train_subset, val_subset = random_split(full_dataset, [train_size, val_size])

# Custom Dataset wrapper to apply transforms
class SubsetWrapper(Dataset):
    def __init__(self, subset, transform=None):
        self.subset = subset
        self.transform = transform

    def __getitem__(self, index):
        x, y = self.subset[index]
        if self.transform:
            x = self.transform(x)
        return x, y

    def __len__(self):
        return len(self.subset)

# Apply different transforms to train and validation datasets
train_dataset = SubsetWrapper(train_subset, transform=train_transform)
val_dataset = SubsetWrapper(val_subset, transform=val_test_transform)

BATCH_SIZE = 32
NUM_WORKERS = os.cpu_count() if device == "cpu" else torch.cuda.device_count()

train_loader = DataLoader(
    dataset=train_dataset,
    batch_size=BATCH_SIZE,
    shuffle=True,
    num_workers=4
)

val_loader = DataLoader(
    dataset=val_dataset,
    batch_size=BATCH_SIZE,
    shuffle=False,
    num_workers=4
)

def display_images_from_dataloader(dataloader, num_images=8):
    mean = torch.tensor([0.485, 0.456, 0.406])
    std = torch.tensor([0.229, 0.224, 0.225])
    
    # 从传入的dataloader提取下一个批次的数据
    images, labels = next(iter(dataloader))
    num_images = min(num_images, images.size(0))
    
    plt.figure(figsize=(4 * num_images, 10))

    for i in range(num_images):
        plt.subplot(1, num_images, i + 1)
        img = images[i] * std.view(3, 1, 1) + mean.view(3, 1, 1)    #逆归一化
        img = img.numpy().transpose((1, 2, 0))    #C,H,W -> H,W,C
        
        label_idx = labels[i].item()
        # category_name = cat_names_list[label_idx]
        
        plt.imshow(img)
        # plt.title(category_name)
        plt.axis('off')
        plt.grid(False)
        plt.tight_layout()
        
    plt.show()
    
# if __name__ == '__main__':
#     torch.manual_seed(21)
#     display_images_from_dataloader(train_loader)
